#!/usr/bin/env python
"""
:mod:`pynchon.bio.profile`
==========================

Provides array objects to store N-dimensional profiles.
"""
from itertools import product
from pynchon.core.point import PointVector0dFloat64, PointArray0dFloat64
from pynchon.bio.data.alphabet import NUMPY_PROXY, IUPAC_GAP
from pynchon.core.symbol import *
from pynchon.bio.aa import AA
from pynchon.bio.nt import Nt, Codon


def to_counts(array, elements, proxy =None):
    """
    Counts the number of occurrences of elements in the "array" for each vector
    along the last axis separately. Returns an array of the same rank as the 
    input, but the last axis has the length equal to the number of elements. 
    The counts of "elements" in the array are summed along the last dimension. 
    The input array should be of "S" data type. The "elements" should be 
    compatible with the data type of "array".
    
    Arguments:
    
        - array (``numpy.ndarray`` instance) An array of rank-1 or greater.
        - elements (sequence) Sequence of objects of a compatible data type 
          with "array". 
        - proxy (``dict``) [default: ``NUMPY_PROXY``] see: 
          ``pynchon.core.bio.data.alphabet._Proxy``.
    """
    if not elements:
        raise ValueError('Required argument "elements" not specified or empty.')
    proxy = proxy or NUMPY_PROXY
    shape = array.shape[:-1] + (len(elements,),)
    freq = np.ndarray(shape=shape, dtype=Freq.data_type)
    for i, element in enumerate(elements):
        freq[...,i] = np.sum(array == proxy[element], axis=-1)
    return freq

def to_joint_counts(array, elements):
    """
    Counts the number of occurrences of joint "elements" in array. A joint 
    element is constructed from a column of "array". All possible joint elements
    are the cross-product of N sets of "elements", where N is the number of rows
    in the array. A joint element has length N.
    
    Arguments:
    
        - array (``numpy.ndarray`` instance) An array of rank-2. Joint counts 
          are calculated column-wise.
        - elements (sequence) Sequence of objects of a compatible data type 
          with "array".
    """
    if not elements:
        raise ValueError('Required metadata "elements" not specified or empty.')    
    if len(array.shape) != 2:
        raise ValueError('The array should be of rank-2.')
    elements_num = len(elements)
    freq_rank = array.shape[0]
    elements_crs = tuple(product(*(elements,)*freq_rank))
    freq_shape = (elements_num,) * freq_rank
    freq_size = elements_num ** freq_rank
    freq = np.zeros(shape=(freq_size,), dtype=Freq.data_type)
    for column in array.transpose():
        element = tuple(column.tolist()) # get strings in tuples
        try:
            element_idx = elements_crs.index(element)
            freq[element_idx] +=1
        except ValueError:
            pass # not in elements.
    return freq.reshape(freq_shape)
    
        
class Freq(PointVector0dFloat64):
    """
    Represents a frequency vector i.e. a vectors of counts of elements. The 
    length of the vector is defined by the "elements" item in the "metadata" 
    ``dict``. 
    """
    
    @classmethod
    def _get_elements(cls, metadata):
        elements = metadata.get('elements', ())
        alphabet = metadata.get('alphabet', ())
        if not elements:
            raise ValueError('Metadata item "elements" not specified or empty.')
        if alphabet:
            for element in elements:
                for letter in element:
                    if letter not in alphabet:
                        raise ValueError('Symbol "%s" not in alphabet' % (letter,))
        
        return elements
    
    @classmethod
    def _get_shape(cls, rank_or_shape, elements):
        el_num = len(elements)
        if issubclass(cls, JointFreq):
            if isinstance(rank_or_shape, tuple):
                rank = len(rank_or_shape)
            elif isinstance(rank_or_shape, int):
                rank = rank_or_shape
            else:
                raise ValueError('The argument "rank" should be an int.')
            if not 1 < rank < 5:
                raise ValueError('The argument "rank" is required and should. '
                                 'be within the range <2, 4>')
            shape = (el_num,)*rank
        elif issubclass(cls, Prof):
            shape = list(rank_or_shape)
            if len(shape) != 2:
                raise ValueError('The input array is not of rank-2.')
            shape[-1] = el_num
        elif issubclass(cls, Freq):
            shape = (el_num,)
        return tuple(shape)
    
    @classmethod
    def _get_dims_dt_fval(cls, array):
        dims = len(array.dtype)
        if np.issubdtype(array.dtype, 'S'):
                dt = 'S'
        elif np.issubdtype(array.dtype, 'U'):
                dt = 'U'              
        else:
            value = array.dtype.fields.values()[0]
            dt = value[0].char
        fval = getattr(array, 'fval', None)
        return dims, dt, fval    

    @classmethod
    def _is_numeric_array(cls, array):
        if np.issubdtype(array.dtype, 'float') or \
           np.issubdtype(array.dtype, 'int'):
            return True
        else:
            return False
        
    @classmethod
    def _validate_numeric(cls, array, shape):
        if array.shape != shape:
            raise ValueError("Wrong shape, should be: %s." % (shape,))
        
    @classmethod
    def _validate_symbolic(cls, array):
        if not np.issubdtype(array.dtype, 'S') and not \
            np.issubdtype(array.dtype, 'U'):
            try: 
                chars = set([i[1][0].char for i in array.dtype.fields.items()])
                if len(chars) != 1:
                    raise ValueError('Structured array not homogeneous.')                            
            except AttributeError:
                raise ValueError('Expected array of floats, integers or symbols.')
        
        rank = len(array.shape)
        if not (cls.rank == -1) and not rank == cls.rank:
            raise ValueError('Expected array of rank: %s got: %s' % \
                             (cls.rank, rank))
        
    @classmethod
    def raw_scratch(cls, rank, dtype, fval, metadata, fast, kwargs):
        if 'shape' in kwargs:
            raise ValueError('Got invalid "shape" argument.')
        if 'size' in kwargs:
            raise ValueError('Got invalid "shape" argument.')
        elements = cls._get_elements(metadata)
        shape = cls._get_shape(rank, elements)
        kwargs['shape'] = shape
        return kwargs

    @classmethod
    def raw_array(cls, rank, dtype, fval, metadata, fast, kwargs, array):
        if 'size' in kwargs:
            raise ValueError('Got invalid "size" argument.')        
        elements = cls._get_elements(metadata)
        shape = cls._get_shape(array.shape, elements)
        isnumeric = cls._is_numeric_array(array)
        if isnumeric:
            cls._validate_numeric(array, shape)
            array = np.asarray(array, dtype=cls.data_type)
        else:
            cls._validate_symbolic(array)
            dims, dt, fval = cls._get_dims_dt_fval(array)
            base = [None, 'SymbolVector', 'SymbolArray', 'SymbolCube', \
                          'SymbolHyperCube'][len(array.shape)]
                          
            Symbol = eval('%s%sd%s1' % (base, dims, dt))
            symbol = Symbol(array, fval=fval, metadata=metadata, fast=fast, copy=False)
            if issubclass(cls, JointFreq):
                array = to_joint_counts(symbol, elements)
            else:
                array = to_counts(symbol, elements, metadata.get('proxy', None))   
        return array
            
    @classmethod
    def raw_sequence(cls, rank, dtype, fval, metadata, fast, kwargs, entities):
        elements = cls._get_elements(metadata)
        if 'size' in kwargs:
            raise ValueError('Got invalid "size" argument.')
        element = elements[0]
        entities = "".join(entities)
        if not isinstance(entities, basestring):
            raise ValueError('Expected "basestring" or a sequence of symbols.')        
        dt = 'S' if isinstance(entities, str) else 'U'
        dims = len(element) if isinstance(element, tuple) else 0 
        SymbolVector = eval('SymbolVector%sd%s1' % (dims, dt))
        sv = SymbolVector(entities, metadata=metadata, fast=fast, copy=False)
        freq = to_counts(sv, elements, metadata.get('proxy', None))
        return [freq] # same rank

    def normalize(self):
        """
        """
        return self / np.sum(self, axis=-1)[None].transpose()
  
           
class AAFreq(Freq):
    """
    """
    metadata = AA.metadata

    
class NtFreq(Freq):
    """
    """
    metadata = Nt.metadata   

 
class CodonFreq(Freq):
    """
    """
    metadata = Codon.metadata 


class GapFreq(Freq):
    """
    """
    metadata = (('elements', IUPAC_GAP),)
    


class JointFreq(PointArray0dFloat64, Freq):
    """
    Represents a joint probability distribution function.
    """
    rank = -1
    
    @classmethod
    def raw_sequence(cls, rank, dtype, fval, metadata, fast, kwargs, entities):
        elements = cls._get_elements(metadata)
        if 'size' in kwargs:
            raise ValueError('Got invalid "size" argument.')
        if isinstance(entities, basestring):
            raise ValueError('Expected a valid SymbolArray input.')
        element = elements[0]
        dt = 'S' if isinstance(element[0], str) else 'U'
        dims = len(element) if isinstance(element, tuple) else 0
        SymbolArray = eval('SymbolArray%sd%s1' % (dims, dt))
        sa = SymbolArray(entities, metadata=metadata, copy=False)
        freq = to_joint_counts(sa, elements)
        return [freq] # rank -1

    def normalize(self):
        """
        """
        return self / np.sum(self)


class AAJointFreq(JointFreq, AAFreq):
    """
    """
  
    
class NtJointFreq(JointFreq, NtFreq):
    """
    """

    
class CodonJointFreq(JointFreq, CodonFreq):
    """
    """

            
class Prof(PointArray0dFloat64, Freq):
    """
    """
    @classmethod
    def raw_scratch(cls, rank, dtype, fval, metadata, fast, kwargs):
        """
        """
        if 'shape' in kwargs:
            raise ValueError('Got invalid "shape" argument.')
        if not 'size' in kwargs or not isinstance(kwargs['size'], int):
            raise ValueError('The argument "size" is invalid or missing.')
        elements = cls._get_elements(metadata)
        kwargs['shape'] = cls._get_shape((kwargs.pop('size'), None), elements)
        return kwargs      
     
    @classmethod
    def raw_sequence(cls, rank, dtype, fval, metadata, fast, kwargs, entities):
        return Entity_sequencer(cls, dtype, fval, metadata, entities)
    
   
class AAProf(Prof, AAFreq):
    """
    """
    
    
class NtProf(Prof, NtFreq):
    """
    """


class CodonProf(Prof, CodonFreq):
    """
    """
    
